module grapes_reader_mod

  use eccodes
  use datetime
  use flogger
  use string
  use const_mod
  use latlon_mesh_mod
  use latlon_process_mod
  use latlon_parallel_mod

  implicit none

  character(10), parameter :: var_names(12) = [ &
    'orog      ' , & ! Surface geopotential height
    '2t        ' , & ! Temperature at 2m height
    '10u       ' , & ! U wind component at 10m height
    '10v       ' , & ! V wind component at 10m height
    'sp        ' , & ! Surface pressure
    'msl       ' , & ! Mean sea level pressure
    'gh        ' , & ! Geopotential height
    't         ' , & ! Temperature
    'u         ' , & ! U wind component
    'v         ' , & ! V wind component
    'q         ' , & ! Specific humidity
    'r         '   & ! Relative humidity
  ]

  type grapes_reader_type
    character(256) file_path
  contains
    procedure :: init => grapes_reader_init
    procedure :: open => grapes_reader_open
    procedure :: get_time => grapes_reader_get_time
    procedure, private :: grapes_reader_get_var_2d
    procedure, private :: grapes_reader_get_var_3d
    generic :: get_var => grapes_reader_get_var_2d, &
                          grapes_reader_get_var_3d
    final :: grapes_reader_final
  end type grapes_reader_type

contains

  subroutine grapes_reader_init(this, file_path)

    class(grapes_reader_type), intent(inout) :: this
    character(*), intent(in) :: file_path

    integer ifile, igrib, iret
    integer num_lon, num_lat, num_plev, i, j
    real(r8) start_lon, end_lon, dlon
    real(r8) start_lat, end_lat, dlat
    real(r8), allocatable :: lon(:)
    real(r8), allocatable :: lat(:)
    real(r8) plev(1000) ! Here we use a large enough array to store pressure levels.
    integer idx
    character(30) shortName, typeOfLevel

    this%file_path = file_path

    call codes_open_file(ifile, file_path, 'r')

    ! Get vertical level information (assume pressure levels).
    num_plev = 0
    do
      call codes_grib_new_from_file(ifile, igrib, iret)
      if (iret == CODES_END_OF_FILE) then
        if (num_plev == 0) then
          print *, 'There is no 3D message in ' // trim(file_path) // '!'
          stop 1
        end if
        exit
      end if
      call codes_get(igrib, 'shortName', shortName)
      if (shortName /= 'u') then
        call codes_release(igrib)
        cycle
      end if
      num_plev = num_plev + 1
      call codes_get(igrib, 'typeOfLevel', typeOfLevel)
      if (typeOfLevel /= 'isobaricInhPa') then
        print *, 'The typeOfLevel is not isobaricInhPa!'
        print *, trim(typeOfLevel)
        stop 1
      end if
      call codes_get(igrib, 'level', plev(num_plev))
      call codes_release(igrib)
    end do
    if (proc%is_root_proc()) then
      call log_notice('There are ' // to_str(num_plev) // ' pressure levels in ' // trim(file_path) // '.')
    end if

    call codes_close_file(ifile)

    ! Get horizonal mesh information.
    call codes_open_file(ifile, file_path, 'r')

    call codes_grib_new_from_file(ifile, igrib, iret)
    call codes_get(igrib, 'Ni', num_lon)
    call codes_get(igrib, 'Nj', num_lat)

    allocate(lon(num_lon))
    allocate(lat(num_lat))

    call codes_get(igrib, 'longitudeOfFirstGridPoint', start_lon); start_lon = start_lon * 1e-6 * rad
    call codes_get(igrib, 'longitudeOfLastGridPoint' , end_lon  ); end_lon   = end_lon   * 1e-6 * rad
    call codes_get(igrib, 'latitudeOfFirstGridPoint' , start_lat); start_lat = start_lat * 1e-6 * rad
    call codes_get(igrib, 'latitudeOfLastGridPoint'  , end_lat  ); end_lat   = end_lat   * 1e-6 * rad
    call codes_get(igrib, 'iDirectionIncrement'      , dlon     ); dlon      = dlon      * 1e-6 * rad
    call codes_get(igrib, 'jDirectionIncrement'      , dlat     ); dlat      = dlat      * 1e-6 * rad
    do i = 1, num_lon
      lon(i) = start_lon + (i - 1) * dlon
    end do
    ! NOTE: Here we like from South Pole to North Pole.
    do j = 1, num_lat
      lat(j) = end_lat + (j - 1) * dlat
    end do
    call codes_release(igrib)

    call codes_close_file(ifile)

    call proc%global_mesh%init_global(lon, lat, plev(1:num_plev))

  end subroutine grapes_reader_init

  subroutine grapes_reader_open(this, file_path)

    class(grapes_reader_type), intent(inout) :: this
    character(*), intent(in) :: file_path

    logical exist

    this%file_path = file_path

    inquire(file=file_path, exist=exist)
    if (.not. exist) then
      if (proc%is_root_proc()) then
        call log_error('File ' // trim(file_path) // ' does not exist!')
      end if
      stop 1
    end if

  end subroutine grapes_reader_open

  type(datetime_type) function grapes_reader_get_time(this) result(res)

    class(grapes_reader_type), intent(in) :: this

    integer ifile, igrib, iret
    character(10) date, time

    call codes_open_file(ifile, this%file_path, 'r')

    call codes_grib_new_from_file(ifile, igrib, iret)
    call codes_get(igrib, 'date', date)
    call codes_get(igrib, 'time', time)
    call codes_release(igrib)
    call codes_close_file(ifile)

    res = create_datetime(trim(date) // trim(time), '%Y%m%d%H%M')

  end function grapes_reader_get_time

  subroutine grapes_reader_get_var_2d(this, short_name, array)

    class(grapes_reader_type), intent(in) :: this
    character(*), intent(in) :: short_name
    real(r8), intent(out) :: array(proc%local_mesh%lon_lb:proc%local_mesh%lon_ub, &
                                   proc%local_mesh%lat_lb:proc%local_mesh%lat_ub)

    integer ifile, igrib, ierr, i, j
    integer numberOfPoints
    character(30) shortName, typeOfLevel
    real(r8), allocatable :: global_lons_dummy(:)
    real(r8), allocatable :: global_lats_dummy(:)
    real(r8), allocatable :: global_array_dummpy(:)
    real(r8), allocatable :: global_array(:,:)

    ! NOTE: Each process open GRIB file for taking its own part of array.
    ! Silly ecCodes library does not provide an interface to read part of the array!
    call codes_open_file(ifile, this%file_path, 'r')

    do
      call codes_grib_new_from_file(ifile, igrib, ierr)
      if (ierr == CODES_END_OF_FILE) exit

      call codes_get(igrib, 'shortName', shortName)
      if (shortName /= short_name) then
        call codes_release(igrib)
        cycle
      end if

      call codes_get(igrib, 'numberOfPoints', numberOfPoints)
      ! Here we only read one pressure level.
      if (.not. allocated(global_array)) then
        allocate(global_lons_dummy(numberOfPoints))
        allocate(global_lats_dummy(numberOfPoints))
        allocate(global_array_dummpy(numberOfPoints))
        allocate(global_array(proc%global_mesh%num_lon,proc%global_mesh%num_lat))
      end if
      call codes_grib_get_data_real8(igrib, global_lats_dummy, global_lons_dummy, global_array_dummpy)
      global_array = reshape(global_array_dummpy, [proc%global_mesh%num_lon,proc%global_mesh%num_lat])

      ! Reverse meridional direction to make South Pole as the starting point.
      do j = proc%local_mesh%lat_ibeg, proc%local_mesh%lat_iend
        do i = proc%local_mesh%lon_ibeg, proc%local_mesh%lon_iend
          array(i,j) = global_array(i,proc%global_mesh%num_lat-j+1)
        end do
      end do
      call codes_release(igrib)
    end do

    call codes_close_file(ifile)

    if (.not. allocated(global_array)) then
      call log_error('Failed to get variable ' // trim(short_name) // '!')
    else
      deallocate(global_lons_dummy)
      deallocate(global_lats_dummy)
      deallocate(global_array_dummpy)
      deallocate(global_array)
    end if

    call fill_halo(array)

  end subroutine grapes_reader_get_var_2d

  subroutine grapes_reader_get_var_3d(this, short_name, array)

    class(grapes_reader_type), intent(in) :: this
    character(*), intent(in) :: short_name
    real(r8), intent(out) :: array(proc%local_mesh%lon_lb:proc%local_mesh%lon_ub, &
                                   proc%local_mesh%lat_lb:proc%local_mesh%lat_ub, &
                                   proc%local_mesh%lev_lb:proc%local_mesh%lev_ub)

    integer ifile, igrib, ierr, i, j, k
    integer numberOfPoints
    character(30) shortName, typeOfLevel
    real(r8) level
    real(r8), allocatable :: global_lons_dummy(:)
    real(r8), allocatable :: global_lats_dummy(:)
    real(r8), allocatable :: global_array_dummpy(:)
    real(r8), allocatable :: global_array(:,:)

    ! NOTE: Each process open GRIB file for taking its own part of array.
    ! Silly ecCodes library does not provide an interface to read part of the array!
    call codes_open_file(ifile, this%file_path, 'r')

    k = 0
    do
      call codes_grib_new_from_file(ifile, igrib, ierr)
      if (ierr == CODES_END_OF_FILE) exit

      call codes_get(igrib, 'shortName', shortName)
      if (shortName /= short_name) then
        call codes_release(igrib)
        cycle
      end if

      call codes_get(igrib, 'numberOfPoints', numberOfPoints)
      call codes_get(igrib, 'typeOfLevel', typeOfLevel)
      call codes_get(igrib, 'level', level)
      ! Here we only read one pressure level.
      if (.not. allocated(global_array)) then
        allocate(global_lons_dummy(numberOfPoints))
        allocate(global_lats_dummy(numberOfPoints))
        allocate(global_array_dummpy(numberOfPoints))
        allocate(global_array(proc%global_mesh%num_lon,proc%global_mesh%num_lat))
      end if
      call codes_grib_get_data_real8(igrib, global_lats_dummy, global_lons_dummy, global_array_dummpy)
      global_array = reshape(global_array_dummpy, [proc%global_mesh%num_lon,proc%global_mesh%num_lat])
      k = k + 1
      ! Reverse meridional direction to make South Pole as the starting point.
      do j = proc%local_mesh%lat_ibeg, proc%local_mesh%lat_iend
        do i = proc%local_mesh%lon_ibeg, proc%local_mesh%lon_iend
          array(i,j,k) = global_array(i,proc%global_mesh%num_lat-j+1)
        end do
      end do
      call codes_release(igrib)
    end do

    call codes_close_file(ifile)

    deallocate(global_lons_dummy)
    deallocate(global_lats_dummy)
    deallocate(global_array_dummpy)
    deallocate(global_array)

    call fill_halo(array)

  end subroutine grapes_reader_get_var_3d

  subroutine grapes_reader_final(this)

    type(grapes_reader_type), intent(inout) :: this

  end subroutine grapes_reader_final

end module grapes_reader_mod
